#include "oci_builder.h"
#include "utils.h"
#include <stdio.h>
#include <stdlib.h>
#include <json-c/json.h>
#include <sys/stat.h>

// 将文件添加到blobs目录, 并返回其digest和大小
static int add_file_to_blobs(const char* file_path, const char* output_dir, char* digest_out, long* size_out) {
    FILE* fp = fopen(file_path, "rb");
    if (!fp) {
        LOG_ERROR("无法打开文件 %s", file_path);
        return -1;
    }

    fseek(fp, 0, SEEK_END);
    *size_out = ftell(fp);
    fseek(fp, 0, SEEK_SET);

    unsigned char* buffer = malloc(*size_out);
    if (!buffer) {
        LOG_ERROR("内存分配失败");
        fclose(fp);
        return -1;
    }
    
    if (fread(buffer, 1, *size_out, fp) != *size_out) {
        LOG_ERROR("读取文件失败: %s", file_path);
        free(buffer);
        fclose(fp);
        return -1;
    }
    fclose(fp);

    char hash_hex[SHA256_HEX_LEN + 1];
    calculate_sha256_hex(buffer, *size_out, hash_hex);
    
    // <--- 修改点: 缓冲区大小增加1, 确保能容纳 \0
    snprintf(digest_out, SHA256_HEX_LEN + 7 + 1, "sha256:%s", hash_hex);

    char blob_path[1024];
    snprintf(blob_path, sizeof(blob_path), "%s/blobs/sha256/%s", output_dir, hash_hex);
    FILE* blob_fp = fopen(blob_path, "wb");
    if (!blob_fp) {
        LOG_ERROR("无法创建 blob 文件: %s", blob_path);
        free(buffer);
        return -1;
    }
    fwrite(buffer, 1, *size_out, blob_fp);
    fclose(blob_fp);
    free(buffer);

    return 0;
}

static json_object* create_oci_image_manifest(const char* config_digest, long config_size,
                                             const char* lba_map_digest, long lba_map_size) {
    json_object *manifest = json_object_new_object();
    json_object_object_add(manifest, "schemaVersion", json_object_new_int(2));

    // Config Layer
    json_object *config_layer = json_object_new_object();
    json_object_object_add(config_layer, "mediaType", json_object_new_string("application/vnd.oci.image.config.v1+json"));
    json_object_object_add(config_layer, "digest", json_object_new_string(config_digest));
    json_object_object_add(config_layer, "size", json_object_new_int64(config_size));
    json_object_object_add(manifest, "config", config_layer);

    // Layers Array
    json_object *layers_array = json_object_new_array();

    // LBA Map Layer (Custom)
    json_object *lba_layer = json_object_new_object();
    json_object_object_add(lba_layer, "mediaType", json_object_new_string("application/vnd.ocis.layer.v1.lba-map"));
    json_object_object_add(lba_layer, "digest", json_object_new_string(lba_map_digest));
    json_object_object_add(lba_layer, "size", json_object_new_int64(lba_map_size));
    // Annotation to hint containerd
    json_object *annotations = json_object_new_object();
    json_object_object_add(annotations, "containerd.io/snapshotter", json_object_new_string("ublk"));
    json_object_object_add(lba_layer, "annotations", annotations);

    json_object_array_add(layers_array, lba_layer);
    json_object_object_add(manifest, "layers", layers_array);

    return manifest;
}

static int create_oci_index(const char* output_dir, const char* manifest_digest, long manifest_size) {
    json_object *index = json_object_new_object();
    json_object_object_add(index, "schemaVersion", json_object_new_int(2));

    json_object *manifests_array = json_object_new_array();
    json_object *manifest_entry = json_object_new_object();
    json_object_object_add(manifest_entry, "mediaType", json_object_new_string("application/vnd.oci.image.manifest.v1+json"));
    json_object_object_add(manifest_entry, "digest", json_object_new_string(manifest_digest));
    json_object_object_add(manifest_entry, "size", json_object_new_int64(manifest_size));

    json_object *platform = json_object_new_object();
    json_object_object_add(platform, "architecture", json_object_new_string("amd64"));
    json_object_object_add(platform, "os", json_object_new_string("linux"));
    json_object_object_add(manifest_entry, "platform", platform);

    json_object_array_add(manifests_array, manifest_entry);
    json_object_object_add(index, "manifests", manifests_array);

    char index_path[1024];
    snprintf(index_path, sizeof(index_path), "%s/index.json", output_dir);
    if (json_object_to_file_ext(index_path, index, JSON_C_TO_STRING_PRETTY) != 0) {
        LOG_ERROR("写入 index.json 文件失败");
        json_object_put(index);
        return -1;
    }
    json_object_put(index);
    return 0;
}


int build_oci_layout(const char* output_dir) {
    LOG_INFO("阶段 3: 构建 OCI 兼容的镜像布局...");

    char lba_map_path[1024];
    snprintf(lba_map_path, sizeof(lba_map_path), "%s/lba_map.bin", output_dir);

    // <--- 修改点: 缓冲区大小增加1
    char lba_map_digest[SHA256_HEX_LEN + 7 + 1];
    long lba_map_size;
    if (add_file_to_blobs(lba_map_path, output_dir, lba_map_digest, &lba_map_size) != 0) {
        LOG_ERROR("无法将 lba_map.bin 添加到 blobs");
        return -1;
    }
    LOG_INFO("LBA Map Blob: digest=%s, size=%ld", lba_map_digest, lba_map_size);

    char config_path[1024];
    // <--- 修改点: 缓冲区大小增加1
    char config_digest[SHA256_HEX_LEN + 7 + 1];
    long config_size;
    snprintf(config_path, sizeof(config_path), "%s/config.json.tmp", output_dir);
    FILE* cfg_fp = fopen(config_path, "w");
    fprintf(cfg_fp, "{}");
    fclose(cfg_fp);
    if (add_file_to_blobs(config_path, output_dir, config_digest, &config_size) != 0) {
        LOG_ERROR("创建 config blob 失败");
        remove(config_path);
        return -1;
    }
    remove(config_path);
    LOG_INFO("Config Blob: digest=%s, size=%ld", config_digest, config_size);

    json_object* image_manifest = create_oci_image_manifest(config_digest, config_size, lba_map_digest, lba_map_size);
    char image_manifest_path[1024];
    // <--- 修改点: 缓冲区大小增加1
    char image_manifest_digest[SHA256_HEX_LEN + 7 + 1];
    long image_manifest_size;
    snprintf(image_manifest_path, sizeof(image_manifest_path), "%s/image_manifest.json.tmp", output_dir);
    json_object_to_file_ext(image_manifest_path, image_manifest, JSON_C_TO_STRING_PRETTY);
    json_object_put(image_manifest);
    if (add_file_to_blobs(image_manifest_path, output_dir, image_manifest_digest, &image_manifest_size) != 0) {
        LOG_ERROR("创建 image manifest blob 失败");
        remove(image_manifest_path);
        return -1;
    }
    remove(image_manifest_path);
    LOG_INFO("Image Manifest Blob: digest=%s, size=%ld", image_manifest_digest, image_manifest_size);

    if (create_oci_index(output_dir, image_manifest_digest, image_manifest_size) != 0) return -1;
    LOG_INFO("index.json 创建成功");

    char oci_layout_path[1024];
    snprintf(oci_layout_path, sizeof(oci_layout_path), "%s/oci-layout", output_dir);
    FILE* oci_layout_fp = fopen(oci_layout_path, "w");
    fprintf(oci_layout_fp, "{\"imageLayoutVersion\": \"1.0.0\"}");
    fclose(oci_layout_fp);
    LOG_INFO("oci-layout 文件创建成功");

    LOG_INFO("阶段 3 完成: OCI 布局构建完毕。");
    return 0;
}